# Author: Hibari
# 2025年11月11日23时57分11秒
# syh19990131@gmail.com

import pandas as pd
import json
import re
from pathlib import Path


def extract_vegetation_data_from_excel(file_path, sheet_name='总植被'):
    """从Excel文件中提取植被数据"""
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path, sheet_name=sheet_name)

        results = []

        for idx, row in df.iterrows():
            # 跳过表头行
            if pd.isna(row.iloc[0]) or row.iloc[0] in ['时期', '']:
                continue

            # 获取时期/朝代
            period = row.iloc[0] if not pd.isna(row.iloc[0]) else '未知时期'

            # 获取植被描述原文
            description = row.iloc[1] if len(row) > 1 and not pd.isna(row.iloc[1]) else None

            # 获取出处
            source = row.iloc[2] if len(row) > 2 and not pd.isna(row.iloc[2]) else None

            # 获取植被分布
            distribution = row.iloc[3] if len(row) > 3 and not pd.isna(row.iloc[3]) else None

            # 获取植被类型
            vegetation_type = row.iloc[4] if len(row) > 4 and not pd.isna(row.iloc[4]) else None

            # 获取植被疏密情况
            density = row.iloc[5] if len(row) > 5 and not pd.isna(row.iloc[5]) else None

            # 创建数据记录
            record = {
                'period': str(period),
                'description': description,
                'source': source,
                'distribution': distribution,
                'vegetation_type': vegetation_type,
                'density': density
            }

            # 只添加有描述内容的记录
            if description:
                results.append(record)

        return results

    except Exception as e:
        print(f"读取文件时出错: {e}")
        return []


def extract_all_dynasty_sheets(file_path):
    """提取所有朝代sheet的数据"""
    all_data = []

    try:
        # 获取所有sheet名称
        excel_file = pd.ExcelFile(file_path)
        sheet_names = excel_file.sheet_names

        for sheet_name in sheet_names:
            if sheet_name not in ['总植被']:  # 跳过总植被sheet，因为我们已经单独处理
                print(f"处理sheet: {sheet_name}")
                sheet_data = extract_vegetation_data_from_excel(file_path, sheet_name)
                all_data.extend(sheet_data)

    except Exception as e:
        print(f"处理各朝代sheet时出错: {e}")

    return all_data


def main():
    # 输入文件路径
    input_file = './data/04植被 - 总数据和各朝代数据.xlsx'

    # 输出文件路径
    output_file = './outputs/vegetation_data.json'

    print("开始提取植被数据...")

    # 提取总植被sheet的数据
    total_vegetation_data = extract_vegetation_data_from_excel(input_file, sheet_name='总植被')
    print(f"从'总植被' sheet提取了 {len(total_vegetation_data)} 条记录")

    # 提取各朝代sheet的数据
    dynasty_data = extract_all_dynasty_sheets(input_file)
    print(f"从各朝代sheet提取了 {len(dynasty_data)} 条记录")

    # 合并所有数据
    all_data = total_vegetation_data + dynasty_data

    # 保存为JSON文件
    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(all_data, f, ensure_ascii=False, indent=2)

    print(f"数据已保存到: {output_file}")

    # 显示前几条数据作为示例
    print("\n前5条数据示例:")
    for i, record in enumerate(all_data[:5], 1):
        print(f"\n记录 {i}:")
        print(f"  时期(period): {record['period']}")
        print(f"  分布(distribution): {record['distribution']}")
        print(f"  植被类型(vegetation_type): {record['vegetation_type']}")
        print(f"  疏密情况(density): {record['density']}")
        if record['description']:
            desc_preview = record['description'][:100] + "..." if len(record['description']) > 100 else record[
                'description']
            print(f"  描述预览: {desc_preview}")

    print(f"\n总共提取了 {len(all_data)} 条有效记录")

    # 统计各时期的数据量
    period_count = {}
    for record in all_data:
        period = record['period']
        period_count[period] = period_count.get(period, 0) + 1

    print("\n各时期数据统计:")
    for period, count in sorted(period_count.items()):
        print(f"  {period}: {count} 条记录")

    return all_data


if __name__ == '__main__':
    main()